import numpy as np
import tensorflow as tf
from tensorflow.keras import layers

class PositionalEncoding(layers.Layer):
  """修正后的位置编码层（支持序列化）"""
  def __init__(self, d_model, max_len=5000, **kwargs):  # 添加**kwargs接收父类参数
      super().__init__(**kwargs)  # 传递父类参数
      self.d_model = d_model
      self.max_len = max_len
      
      # 生成位置编码矩阵
      position = np.arange(max_len)[:, np.newaxis]
      div_term = np.exp(
          np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model)
      )
      pe = np.zeros((max_len, d_model))
      pe[:, 0::2] = np.sin(position * div_term)
      pe[:, 1::2] = np.cos(position * div_term)
      self.pe = tf.constant(pe[np.newaxis, :], dtype=tf.float32)

  def call(self, x):
      return x + self.pe[:, :tf.shape(x)[1], :]

  # 新增序列化配置方法
  def get_config(self):
      config = super().get_config()
      config.update({
          "d_model": self.d_model,
          "max_len": self.max_len
      })
      return config

  # 新增反序列化方法
  @classmethod
  def from_config(cls, config):
      return cls(**config)